package cn.hiapi.core.basic.sharding.algorithm;

import cn.hiapi.core.basic.sharding.utils.ShardingUtils;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.commons.lang3.time.DateUtils;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;

import java.util.*;

/**
 * @author AdinZ
 * @date 2024/1/2 16:04
 */
public class DataMonthShardingAlgorithm implements StandardShardingAlgorithm<Long> {


    private String format;

    @Override
    public void init(Properties props) {
        this.format = props.getProperty("format", "yyyyMM");
    }

    /**
     * Sharding.
     *
     * @param availableTargetNames available data sources or table names
     * @param value                sharding value
     * @return sharding result for data source or table name
     */
    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<Long> value) {
        String table = ShardingUtils.getTableTime(value.getLogicTableName(), value.getValue(), this.format);
        if (availableTargetNames.contains(table)) {
            return table;
        }
        return null;
    }


    /**
     * Sharding.
     *
     * @param availableTargetNames available data sources or table names
     * @param value                sharding value
     * @return sharding results for data sources or table names
     */
    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Long> value) {

        long rangeLowerDate;
        if (value.getValueRange().hasLowerBound()) {
            rangeLowerDate = value.getValueRange().lowerEndpoint();
        } else {
            rangeLowerDate = System.currentTimeMillis();
        }

        // 获取到范围查找的最大值，如果没有配置最大值，设置为当前时间 + 1 月
        // 这里需要注意，你的项目里面这样做是否合理
        long rangeUpperDate;
        if (value.getValueRange().hasUpperBound()) {
            rangeUpperDate = value.getValueRange().upperEndpoint();
        } else {
            // 往后延一个月
            rangeUpperDate = DateUtils.addMonths(new Date(), 1).getTime();
        }
        Set<String> tables = rangeTable(value.getLogicTableName(), rangeLowerDate, rangeUpperDate);
        Set<String> result = new HashSet<>();
        tables.forEach(table -> {
            if (availableTargetNames.contains(table) || ShardingAlgorithmTable.TABLES.contains(table)) {
                result.add(table);
            }
        });

        return result;
    }


    @Override
    public String getType() {
        return "TIME";
    }


    private Set<String> rangeTable(String logicTable, long rangeLowerDate, long rangeUpperDate) {
        Set<String> tableNames = new HashSet<>();
        // 过滤那些存在的表
        while (rangeLowerDate < rangeUpperDate) {
            String actuallyTableName = logicTable + ShardingUtils.getTableTime(logicTable, rangeLowerDate, this.format);
            tableNames.add(actuallyTableName);
            rangeLowerDate = DateUtils.addMonths(new Date(rangeLowerDate), 1).getTime();
        }
        return tableNames;
    }

}
